from datetime import datetime

import openai
from langchain.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
# 加载 .env 到环境变量
from dotenv import load_dotenv, find_dotenv
from langfuse.decorators import observe, langfuse_context

_ = load_dotenv(find_dotenv())

model = ChatOpenAI(model="gpt-3.5-turbo")
prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template("Say hello to {input}!")
    ]
)

# 定义输出解析器
parser = StrOutputParser()
# 定义chain，也就是LCEL
chain = (
        {"input": RunnablePassthrough()}
        | prompt
        | model
        | parser
)


@observe()
def run():
    langfuse_context.update_current_trace(
        name="LangChainDemo",
        user_id="wzr",
    )
    # get the langchain handler for the current trace(获取当前trace中langchain提供的处理器)
    langfuse_handler = langfuse_context.get_current_langchain_handler()
    return chain.invoke(input="AGIClass", config={"callbacks": [langfuse_handler]})


print(run())
